{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Week 1: CUDA operations in PyTorch. Introduction to benchmarking.\n",
        "In this seminar, we'll learn a bit more about the things one needs to keep in mind when using GPU computations, both in general and in PyTorch. We'll also see a couple of examples on benchmarking code in Python; again, there are some caveats with CUDA.\n",
        "\n",
        "First, let's import PyTorch and get some information about the currently used device:"
      ],
      "metadata": {
        "id": "4yYIhuphpTyt"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "089F8Q-yK9-w",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d8c164a3-38a0-493a-e8d1-9cee18428c95"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 1
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "torch.cuda.is_available()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.get_device_properties(0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C9iz9uYOn5JY",
        "outputId": "36d2ef6a-6a4d-4604-f1c5-74e95fe412bd"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "_CudaDeviceProperties(name='Tesla T4', major=7, minor=5, total_memory=15102MB, multi_processor_count=40, uuid=50acfeb7-0f59-b1b1-33e8-b0a8eadc959f, L2_cache_size=4MB)"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Memory allocation\n",
        "As discussed in the lecture, GPU memory is separate from the CPU memory and needs to be explicitly allocated, triggering a host-device synchronization. PyTorch uses a caching memory allocator to repurpose already available but unused memory fragments.\n",
        "\n",
        "See this example: we allocate the memory inside the function scope, so the tensor is deleted as soon as the scope is left:"
      ],
      "metadata": {
        "id": "i3_7snMwx5Yq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def allocate_empty_tensor(dim_size):\n",
        "    a = torch.empty(4096, dim_size, dtype=torch.float32, device=\"cuda\")"
      ],
      "metadata": {
        "id": "Maafa5Nehba-"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "allocate_empty_tensor(2048)"
      ],
      "metadata": {
        "id": "LhI9_OvaoOIi"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Printing the allocated memory size gives us an expected number of zero bytes:"
      ],
      "metadata": {
        "id": "XAxd3nkJ3UsK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.memory_allocated(0)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rez0XGH6gGF9",
        "outputId": "d8c8faa4-d77f-461c-b4fd-0d6eeb345be6"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "However, the GPU memory is still in use by the process; this is shown by `torch.cuda.memory_reserved` or `nvidia-smi` from the terminal. As a result, working in a shared GPU environment can leave a lot of unused yet allocated memory."
      ],
      "metadata": {
        "id": "qzcQ2ohm3d2v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.memory_reserved()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H6IOufbXoQwI",
        "outputId": "41ce037d-3e0e-4ca8-e759-183c49afcabd"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "33554432"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1iM6LZBWIDT1",
        "outputId": "f5ba9d99-e8ee-4511-d7a7-dfa68fb4d92f"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tue Jan 21 10:59:44 2025       \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |\n",
            "|-----------------------------------------+----------------------+----------------------+\n",
            "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                                         |                      |               MIG M. |\n",
            "|=========================================+======================+======================|\n",
            "|   0  Tesla T4                       Off | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   62C    P0              31W /  70W |    135MiB / 15360MiB |      3%      Default |\n",
            "|                                         |                      |                  N/A |\n",
            "+-----------------------------------------+----------------------+----------------------+\n",
            "                                                                                         \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| Processes:                                                                            |\n",
            "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
            "|        ID   ID                                                             Usage      |\n",
            "|=======================================================================================|\n",
            "+---------------------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's clear the cache now and see what happens:"
      ],
      "metadata": {
        "id": "JUK6clp83nGu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "torch.cuda.memory_reserved()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SSPz5VWYopb7",
        "outputId": "efe1ba27-45d1-452c-d2bb-665c93081288"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note that this operation triggers a CPU-GPU synchronization, and thus using it in your code can significantly hurt the performance. It's almost always better to carefully manage the lifetime of your GPU tensors and avoid excessive allocations than to empty the cache.\n",
        "\n",
        "Now, let's see how this cache is reused by allocating two tensors in a row: first a larger one, then a smaller one."
      ],
      "metadata": {
        "id": "m3m-rcF64j1l"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "allocate_empty_tensor(2048)\n",
        "torch.cuda.memory_reserved()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8YHdjvrn32_F",
        "outputId": "3aafb500-e2ac-4557-c356-a5cc1791a836"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "33554432"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "allocate_empty_tensor(1024)\n",
        "torch.cuda.memory_reserved()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4qKRSbS634vg",
        "outputId": "86d5dbde-e3fa-4c73-c849-4b1bd3b18853"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "33554432"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As expected, we reuse the cache, since the array fits into the allocated chunk.\n",
        "\n",
        "However, if we attempt to do this with a larger tensor (3072 elements in the second dimension instead of 2048), we observe something different:"
      ],
      "metadata": {
        "id": "MpLYUugy4hTW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "allocate_empty_tensor(3072)\n",
        "torch.cuda.memory_reserved()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aOe0In1I3_ix",
        "outputId": "4d2ddc31-3de4-4e62-bebe-c02143de81bc"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "83886080"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "What happened? The chunk of memory that was allocated for a 4096x2048 array did not fit a tensor of size 4096x3072: thus, PyTorch needed to allocate an additional segment of a sufficient size while keeping the previous one allocated (in case the user creates a smaller tensor later at some point).\n",
        "\n",
        "In practice, this means that if your code is dealing with tensors of dynamic size (changing batch sizes, sequence lengths or image resolutions), it is recommended to first warm up the cache by allocating tensors for the largest expected input. Otherwise, in the worst case you allocate a quadratic amount of memory with respect to the largest input size instead of a linear one.\n",
        "\n",
        "You can also view more detailed allocation statistics by running `torch.cuda.memory_stats()`"
      ],
      "metadata": {
        "id": "ayR2vJmL4g4a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "memory_stats = torch.cuda.memory_stats()\n",
        "print(memory_stats[\"active.all.allocated\"])\n",
        "print(memory_stats[\"active.all.current\"])\n",
        "print(memory_stats[\"active.all.peak\"])\n",
        "print(memory_stats[\"reserved_bytes.all.current\"])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "leOBOx9OmCqP",
        "outputId": "eeb62032-ea82-4554-c391-96606d6e1175"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4\n",
            "0\n",
            "1\n",
            "83886080\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()\n",
        "print(torch.cuda.memory_stats()[\"reserved_bytes.all.current\"])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BtsyuLdFI6N9",
        "outputId": "8328c8eb-f9dd-4c06-affc-80a5cb95b4ea"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Benchmarking intro\n",
        "The simplest way to check the performance impact of any change is to compare the runtime of code with and without it. Here, we will consider a simple way of doing this in Python.\n",
        "\n",
        "First, let's define two functions that compute a batched version of a dot product for two matrices:"
      ],
      "metadata": {
        "id": "PexpzQZ_xPss"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def batched_dot_mul_sum(a, b):\n",
        "    \"\"\"Computes batched dot by multiplying and summing\"\"\"\n",
        "    return a.mul(b).sum(-1)\n",
        "\n",
        "\n",
        "def batched_dot_bmm(a, b):\n",
        "    \"\"\"Computes batched dot by reducing to bmm\"\"\"\n",
        "    a = a.reshape(-1, 1, a.shape[-1])\n",
        "    b = b.reshape(-1, b.shape[-1], 1)\n",
        "    return torch.bmm(a, b).flatten(-3)\n",
        "\n",
        "\n",
        "# Input for benchmarking\n",
        "x = torch.randn(10000, 64)\n",
        "\n",
        "# Ensure that both functions compute the same output\n",
        "assert batched_dot_mul_sum(x, x).allclose(batched_dot_bmm(x, x))"
      ],
      "metadata": {
        "id": "YH5Wbu_GLLN0"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "To conduct microbenchmarks by running the code hundreds of times and measuring the average execution time, you can use the built-in [timeit](https://docs.python.org/3/library/timeit.html) module. Simply create a Timer object with relevant arguments and call `.timeit()`:"
      ],
      "metadata": {
        "id": "7iJ6zelnyE37"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import timeit\n",
        "\n",
        "t0 = timeit.Timer(\n",
        "    stmt=\"batched_dot_mul_sum(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_mul_sum\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "t1 = timeit.Timer(\n",
        "    stmt=\"batched_dot_bmm(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_bmm\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "print(f\"mul_sum(x, x):  {t0.timeit(100) / 100 * 1e6:>5.1f} us\")\n",
        "print(f\"bmm(x, x):      {t1.timeit(100) / 100 * 1e6:>5.1f} us\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rKAVpQ1hLMxI",
        "outputId": "9ded4113-45b4-44e2-8250-c8131909d1f3"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "mul_sum(x, x):  414.5 us\n",
            "bmm(x, x):      828.6 us\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In IPython, there exist line and cell timeit [magics](https://ipython.readthedocs.io/en/stable/interactive/magics.html#cell-magics). They can be more convenient for smaller cases but allow you a bit less control over the setup."
      ],
      "metadata": {
        "id": "OFm8rggSK3FW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit batched_dot_mul_sum(x, x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NM0hjbs3JxLF",
        "outputId": "d0664017-f8c6-4b9b-edfb-a93a4cf81375"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "356 µs ± 18.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit batched_dot_bmm(x, x)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a28ZIFy_KZl0",
        "outputId": "e9529548-20b4-4f01-bc98-8376833b5dbf"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "871 µs ± 91.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "[torch.utils.benchmark](https://pytorch.org/docs/stable/benchmark_utils.html) copies the API of timeit while helping the user avoid common mistakes.\n",
        "\n",
        "First, let's run it on the same code and data:"
      ],
      "metadata": {
        "id": "F1WqSrVJyEAd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.utils.benchmark as benchmark\n",
        "\n",
        "t0 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_mul_sum(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_mul_sum\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "t1 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_bmm(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_bmm\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "print(t0.timeit(100))\n",
        "print(t1.timeit(100))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GF1XAx1ULSch",
        "outputId": "e476da7f-daca-47fb-ef31-396dee654b94"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a5537bf6bd0>\n",
            "batched_dot_mul_sum(x, x)\n",
            "setup: from __main__ import batched_dot_mul_sum\n",
            "  579.12 us\n",
            "  1 measurement, 100 runs , 1 thread\n",
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a55585655d0>\n",
            "batched_dot_bmm(x, x)\n",
            "setup: from __main__ import batched_dot_bmm\n",
            "  983.38 us\n",
            "  1 measurement, 100 runs , 1 thread\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# in addition, we can set the number of threads for CPU computations\n",
        "num_threads = torch.get_num_threads()\n",
        "print(f\"Benchmarking on {num_threads} threads\")\n",
        "\n",
        "t0 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_mul_sum(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_mul_sum\",\n",
        "    globals={\"x\": x},\n",
        "    num_threads=num_threads,\n",
        "    label=\"Multithreaded batch dot\",\n",
        "    sub_label=\"Implemented using mul and sum\",\n",
        ")\n",
        "\n",
        "t1 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_bmm(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_bmm\",\n",
        "    globals={\"x\": x},\n",
        "    num_threads=num_threads,\n",
        "    label=\"Multithreaded batch dot\",\n",
        "    sub_label=\"Implemented using bmm\",\n",
        ")\n",
        "\n",
        "print(t0.timeit(100))\n",
        "print(t1.timeit(100))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jHfn-xHLLTxm",
        "outputId": "6aefe23a-9064-47f7-b99a-6c26d6f6cfd8"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Benchmarking on 1 threads\n",
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a5537a3d650>\n",
            "Multithreaded batch dot: Implemented using mul and sum\n",
            "setup: from __main__ import batched_dot_mul_sum\n",
            "  556.06 us\n",
            "  1 measurement, 100 runs , 1 thread\n",
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a5537a3d450>\n",
            "Multithreaded batch dot: Implemented using bmm\n",
            "setup: from __main__ import batched_dot_bmm\n",
            "  1.42 ms\n",
            "  1 measurement, 100 runs , 1 thread\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# we can change it globally for PyTorch and measure the impact\n",
        "prev_num_threads = num_threads\n",
        "torch.set_num_threads(2)\n",
        "\n",
        "num_threads = torch.get_num_threads()\n",
        "print(f\"Benchmarking on {num_threads} threads\")\n",
        "\n",
        "t0 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_mul_sum(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_mul_sum\",\n",
        "    globals={\"x\": x},\n",
        "    num_threads=num_threads,\n",
        "    label=\"Multithreaded batch dot\",\n",
        "    sub_label=\"Implemented using mul and sum\",\n",
        ")\n",
        "\n",
        "t1 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_bmm(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_bmm\",\n",
        "    globals={\"x\": x},\n",
        "    num_threads=num_threads,\n",
        "    label=\"Multithreaded batch dot\",\n",
        "    sub_label=\"Implemented using bmm\",\n",
        ")\n",
        "\n",
        "print(t0.timeit(100))\n",
        "print(t1.timeit(100))\n",
        "# in this case, we don't get any speedup, likely due to the overhead\n",
        "\n",
        "torch.set_num_threads(prev_num_threads)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xuHOmnUplY5T",
        "outputId": "2c4cf2c9-5bc7-45c5-e9ff-f06b14af984b"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Benchmarking on 2 threads\n",
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a5537a3c990>\n",
            "Multithreaded batch dot: Implemented using mul and sum\n",
            "setup: from __main__ import batched_dot_mul_sum\n",
            "  758.38 us\n",
            "  1 measurement, 100 runs , 2 threads\n",
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a55586bee10>\n",
            "Multithreaded batch dot: Implemented using bmm\n",
            "setup: from __main__ import batched_dot_bmm\n",
            "  509.14 us\n",
            "  1 measurement, 100 runs , 2 threads\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# by the way, what CPU do we have?\n",
        "!lscpu"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0sbnCapQNirw",
        "outputId": "2f877c54-4cdb-4df3-81d1-4edbfa2ef450"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Architecture:             x86_64\n",
            "  CPU op-mode(s):         32-bit, 64-bit\n",
            "  Address sizes:          46 bits physical, 48 bits virtual\n",
            "  Byte Order:             Little Endian\n",
            "CPU(s):                   2\n",
            "  On-line CPU(s) list:    0,1\n",
            "Vendor ID:                GenuineIntel\n",
            "  Model name:             Intel(R) Xeon(R) CPU @ 2.00GHz\n",
            "    CPU family:           6\n",
            "    Model:                85\n",
            "    Thread(s) per core:   2\n",
            "    Core(s) per socket:   1\n",
            "    Socket(s):            1\n",
            "    Stepping:             3\n",
            "    BogoMIPS:             4000.40\n",
            "    Flags:                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 cl\n",
            "                          flush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc re\n",
            "                          p_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3\n",
            "                           fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand\n",
            "                           hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp \n",
            "                          fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f a\n",
            "                          vx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveop\n",
            "                          t xsavec xgetbv1 xsaves arat md_clear arch_capabilities\n",
            "Virtualization features:  \n",
            "  Hypervisor vendor:      KVM\n",
            "  Virtualization type:    full\n",
            "Caches (sum of all):      \n",
            "  L1d:                    32 KiB (1 instance)\n",
            "  L1i:                    32 KiB (1 instance)\n",
            "  L2:                     1 MiB (1 instance)\n",
            "  L3:                     38.5 MiB (1 instance)\n",
            "NUMA:                     \n",
            "  NUMA node(s):           1\n",
            "  NUMA node0 CPU(s):      0,1\n",
            "Vulnerabilities:          \n",
            "  Gather data sampling:   Not affected\n",
            "  Itlb multihit:          Not affected\n",
            "  L1tf:                   Mitigation; PTE Inversion\n",
            "  Mds:                    Vulnerable; SMT Host state unknown\n",
            "  Meltdown:               Vulnerable\n",
            "  Mmio stale data:        Vulnerable\n",
            "  Reg file data sampling: Not affected\n",
            "  Retbleed:               Vulnerable\n",
            "  Spec rstack overflow:   Not affected\n",
            "  Spec store bypass:      Vulnerable\n",
            "  Spectre v1:             Vulnerable: __user pointer sanitization and usercopy barriers only; no swa\n",
            "                          pgs barriers\n",
            "  Spectre v2:             Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BH\n",
            "                          I: Vulnerable (Syscall hardening enabled)\n",
            "  Srbds:                  Not affected\n",
            "  Tsx async abort:        Vulnerable\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Benchmarking GPU code\n",
        "As we discussed in the lecture, CUDA kernel execution and PyTorch GPU operations are asynchronous.\n",
        "This means that it is possible to launch several kernels in a row before receiving results from the first one. As a consequence, naive benchmarking without synchronization is likely to give you unrealistic results.\n",
        "\n",
        "Let's the same example, but with slightly larger matrices on the GPU:"
      ],
      "metadata": {
        "id": "KHJDeL81zJt4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import timeit\n",
        "\n",
        "x = torch.randn(10000, 1024, device=\"cuda\")\n",
        "\n",
        "t0 = timeit.Timer(\n",
        "    stmt=\"batched_dot_mul_sum(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_mul_sum\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "t1 = timeit.Timer(\n",
        "    stmt=\"batched_dot_bmm(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_bmm\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "# Ran each twice to show difference before/after warmup\n",
        "print(f\"mul_sum(x, x):  {t0.timeit(100) / 100 * 1e6:>5.1f} us\")\n",
        "print(f\"mul_sum(x, x):  {t0.timeit(100) / 100 * 1e6:>5.1f} us\")\n",
        "print(f\"bmm(x, x):      {t1.timeit(100) / 100 * 1e6:>5.1f} us\")\n",
        "print(f\"bmm(x, x):      {t1.timeit(100) / 100 * 1e6:>5.1f} us\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UWPVFvQxLWk9",
        "outputId": "c08a500b-098c-49ef-95c6-9a61f3107a7b"
      },
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "mul_sum(x, x):  733.8 us\n",
            "mul_sum(x, x):   28.2 us\n",
            "bmm(x, x):      1598.6 us\n",
            "bmm(x, x):       29.6 us\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First, we see that the difference between the first and the second runs of timeit is quite noticeable, which happens because of initializing the CUDA context and loading [cuBLAS](https://docs.nvidia.com/cuda/cublas/index.html) — a CUDA library for accelerated linear algebra.\n",
        "\n",
        "Second, the runtimes of two methods seem too small.\n",
        "\n",
        "Let's run the same test with `torch.utils.benchmark` to see a different set of results."
      ],
      "metadata": {
        "id": "xt5nFXZJzn-q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "t0 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_mul_sum(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_mul_sum\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "t1 = benchmark.Timer(\n",
        "    stmt=\"batched_dot_bmm(x, x)\",\n",
        "    setup=\"from __main__ import batched_dot_bmm\",\n",
        "    globals={\"x\": x},\n",
        ")\n",
        "\n",
        "# Run only once since benchmark module does warmup for us\n",
        "print(t0.timeit(100))\n",
        "print(t1.timeit(100))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G2S71HcljicI",
        "outputId": "c5bb9ae5-5d97-42a1-a98f-83ac29a38baa"
      },
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a55586bf110>\n",
            "batched_dot_mul_sum(x, x)\n",
            "setup: from __main__ import batched_dot_mul_sum\n",
            "  498.76 us\n",
            "  1 measurement, 100 runs , 1 thread\n",
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a5537a3f110>\n",
            "batched_dot_bmm(x, x)\n",
            "setup: from __main__ import batched_dot_bmm\n",
            "  158.94 us\n",
            "  1 measurement, 100 runs , 1 thread\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now we have a more realistic set of measurements, which is caused by explicitly triggering the CPU-GPU synchronization and awaiting the results after launching the multiplication. The runtime is comparable to what we had in a previous set of benchmarks, but pay attention that now the matrices being multiplied are 16 times bigger.\n",
        "\n",
        "Let's implement microbenchmarking by ourselves using nothing but [`time.perf_counter()`](https://docs.python.org/3/library/time.html#time.perf_counter) and CUDA methods given by PyTorch. First, we'll run `batched_dot_mul_sum` on the GPU with and without synchronization and see how this affects the results:"
      ],
      "metadata": {
        "id": "k0OG91pQztfY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from time import perf_counter\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "execution_times = []\n",
        "\n",
        "for _ in range(100):\n",
        "    start_time = perf_counter()\n",
        "    batched_dot_mul_sum(x, x)\n",
        "    execution_times.append(perf_counter() - start_time)\n",
        "\n",
        "np.mean(execution_times)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gFyWDoTnrkdw",
        "outputId": "d5700f4f-a2cf-4b26-992e-bb7ef4fa35e0"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5.13540500014642e-05"
            ]
          },
          "metadata": {},
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Compare this with the result that does not compute anything:"
      ],
      "metadata": {
        "id": "daRiA7czQSF1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "execution_times = []\n",
        "\n",
        "for _ in range(100):\n",
        "    start = perf_counter()\n",
        "    execution_times.append(perf_counter() - start)\n",
        "\n",
        "np.mean(execution_times)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bhyGF5-Ftta4",
        "outputId": "6c8705e2-fcfa-47f1-c1d4-93865273e2ff"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "2.1854999999959544e-07"
            ]
          },
          "metadata": {},
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Finally, let's explicitly call `torch.cuda.synchronize` at the end of each iteration to see the difference."
      ],
      "metadata": {
        "id": "KJ8KEHVKQ1Y_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "execution_times = []\n",
        "\n",
        "for _ in range(100):\n",
        "    start_time = perf_counter()\n",
        "    batched_dot_mul_sum(x, x)\n",
        "    torch.cuda.synchronize()\n",
        "    execution_times.append(perf_counter() - start_time)\n",
        "\n",
        "np.mean(execution_times)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kjmFVZdPQjrj",
        "outputId": "f813cc7c-478f-4cc2-f081-636b59d90113"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0005366688899982819"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The same thing applies to the actual models. Let's take `Linear` as the simplest example:"
      ],
      "metadata": {
        "id": "r4AKp3IrW_ei"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x = torch.randn(10000, 512, device=\"cuda\")\n",
        "linear = torch.nn.Linear(512, 1024, bias=False, device=\"cuda\")\n",
        "N_ITERS = 200\n",
        "\n",
        "execution_times = []\n",
        "\n",
        "for _ in range(N_ITERS):\n",
        "    start = perf_counter()\n",
        "    result = linear(x)\n",
        "    execution_times.append(perf_counter() - start)\n",
        "\n",
        "np.mean(execution_times)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E21i4UMYujuQ",
        "outputId": "dfe58690-8cc2-48de-b572-f7b3b9227e11"
      },
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.00018119562499805398"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "execution_times = []\n",
        "\n",
        "for _ in range(N_ITERS):\n",
        "    start = perf_counter()\n",
        "    result = linear(x)\n",
        "    torch.cuda.synchronize()\n",
        "    execution_times.append(perf_counter() - start)\n",
        "\n",
        "np.mean(execution_times)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T6oR06jMvXA3",
        "outputId": "59ce665b-629c-4918-a342-6ee52a7efd21"
      },
      "execution_count": 28,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0030591336149984726"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "execution_times = []\n",
        "\n",
        "start = perf_counter()\n",
        "for _ in range(N_ITERS):\n",
        "    result = linear(x)\n",
        "torch.cuda.synchronize()\n",
        "\n",
        "(perf_counter() - start) / N_ITERS"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bNJ0S1ykwPtZ",
        "outputId": "79e93527-eed0-4350-c066-dc8ebfe5394f"
      },
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0027414358349994927"
            ]
          },
          "metadata": {},
          "execution_count": 29
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Of course, `torch.utils.benchmark` also works:"
      ],
      "metadata": {
        "id": "7C3AJER9XU0t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "t0 = benchmark.Timer(stmt=\"linear(x)\", globals={\"x\": x, \"linear\": linear})\n",
        "print(t0.timeit(100))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ISGCgjslvqXz",
        "outputId": "e69a46b2-ab48-49b4-8477-218e8b2824ae"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<torch.utils.benchmark.utils.common.Measurement object at 0x7a5537a3f110>\n",
            "linear(x)\n",
            "  2.54 ms\n",
            "  1 measurement, 100 runs , 1 thread\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CUDA Streams\n",
        "To execute several operations concurrently, you may use CUDA streams in PyTorch. Below, you can see an example of their usage."
      ],
      "metadata": {
        "id": "j8HSUDAISim3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "cuda = torch.device(\"cuda\")\n",
        "s = torch.cuda.Stream()  # Create a new stream.\n",
        "A = torch.empty((1000, 1000), device=cuda).normal_(0.0, 1.0)\n",
        "with torch.cuda.stream(s):\n",
        "    # sum() may start execution before normal_() finishes!\n",
        "    B = torch.sum(A)"
      ],
      "metadata": {
        "id": "VAwCkCrgmwLw"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "s1 = torch.cuda.Stream()\n",
        "s2 = torch.cuda.Stream()\n",
        "# Initialize cuda tensors here. E.g.:\n",
        "A = torch.rand(1000, 1000, device=\"cuda\")\n",
        "B = torch.rand(1000, 1000, device=\"cuda\")\n",
        "# Wait for the above tensors to initialize.\n",
        "torch.cuda.synchronize()\n",
        "\n",
        "execution_times = []\n",
        "\n",
        "for _ in range(100):\n",
        "    start = perf_counter()\n",
        "    with torch.cuda.stream(s1):\n",
        "        C = torch.mm(A, A)\n",
        "    with torch.cuda.stream(s2):\n",
        "        D = torch.mm(B, B)\n",
        "    # Wait for C and D to be computed.\n",
        "    torch.cuda.synchronize()\n",
        "    execution_times.append(perf_counter() - start)\n",
        "\n",
        "np.mean(execution_times)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fJOJ1bi4nImI",
        "outputId": "2e19fb64-d191-4239-9198-3959ddcb1e32"
      },
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0014858745499930137"
            ]
          },
          "metadata": {},
          "execution_count": 32
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# next, let's compute C and D sequentially\n",
        "\n",
        "execution_times = []\n",
        "\n",
        "for _ in range(100):\n",
        "    start = perf_counter()\n",
        "    C = torch.mm(A, A)\n",
        "    D = torch.mm(B, B)\n",
        "    # Wait for C and D to be computed.\n",
        "    torch.cuda.synchronize()\n",
        "    execution_times.append(perf_counter() - start)\n",
        "\n",
        "np.mean(execution_times)\n",
        "# the speed is even higher in this case"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RRw1SQ9hmNiW",
        "outputId": "b2b86821-038c-4a7a-af58-e5254be80573"
      },
      "execution_count": 33,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.0016013037200036705"
            ]
          },
          "metadata": {},
          "execution_count": 33
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As you can see in this example, the usefulness of streams can be limited: concurrent kernels need to underutilize the GPU and yet take long enough to compute. It might still be a good idea if you are trying to compute some expression and fetch data for the next computation simultaneously (for example, when training examples are large or in case of offloading).\n",
        "The example below demonstrates how to use streams for that purpose."
      ],
      "metadata": {
        "id": "iog1q8KTuYtx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "compute_stream = torch.cuda.Stream()\n",
        "h2d_stream = torch.cuda.Stream()\n",
        "\n",
        "A = torch.rand(10240, 10240, device=\"cuda\")\n",
        "B = [torch.rand(4096, 4096, device=\"cpu\").pin_memory() for _ in range(100)]"
      ],
      "metadata": {
        "id": "3dR-9sTJjbsP"
      },
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def sequential_execution():\n",
        "  torch.matmul(A, A)\n",
        "  for matrix in B:\n",
        "    matrix.to(\"cuda\", non_blocking=True)"
      ],
      "metadata": {
        "id": "5fGIhMiNr5zF"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "sequential_execution()\n",
        "torch.cuda.synchronize()"
      ],
      "metadata": {
        "id": "nW1luh_cj3QZ",
        "outputId": "5706d0df-a2ef-4332-aa7c-1d97dfe4d5f1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.18 s ± 16.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "def stream_execution():\n",
        "  with torch.cuda.stream(compute_stream):\n",
        "    torch.matmul(A, A)\n",
        "  with torch.cuda.stream(h2d_stream):\n",
        "    for matrix in B:\n",
        "      matrix.to(\"cuda\")"
      ],
      "metadata": {
        "id": "2HpXlmGtj98h"
      },
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "stream_execution()\n",
        "torch.cuda.synchronize()"
      ],
      "metadata": {
        "id": "E5S06upZkF93",
        "outputId": "95e3b10d-a188-4188-dc41-9ebfd3dca2d5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "692 ms ± 148 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Debugging asynchronous code\n",
        "Finding sources of errors in GPU-reliant code in PyTorch can also be difficult. Let's see this on a standard example of an off-by-one error and incorrect index for the embedding layer:"
      ],
      "metadata": {
        "id": "TC5UISMiSAfs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile incorrect_index.py\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "embedding = nn.Embedding(1024,32).to('cuda')\n",
        "# 1024 > 1023 (largest index in the created embedding layer)\n",
        "input = torch.full((1,1),1024,dtype=torch.long, device='cuda')\n",
        "\n",
        "# out-of-bounds access\n",
        "embedding_for_index = embedding(input)\n",
        "\n",
        "result = torch.sigmoid(embedding_for_index)\n",
        "loss = result.sum()\n",
        "print(loss.item())\n",
        "print(loss)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "75YMFmOdzhAw",
        "outputId": "04373b69-917a-4716-f812-cab9dbedbfb5"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing incorrect_index.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "If we attempt to run the code as is, we'll see an error after the point in which we trigger the CPU-GPU synchronization (`loss.item`). This gives us no clues about what exactly caused an error:"
      ],
      "metadata": {
        "id": "dxMcIkDmYBFn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!python incorrect_index.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "whEGTI1-VVQ6",
        "outputId": "6a5d2758-68f2-4876-e17b-a01c2e49aaca"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [0,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [1,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [2,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [3,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [4,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [5,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [6,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [7,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [8,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [9,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [10,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [11,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [12,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [13,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [14,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [15,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [16,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [17,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [18,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [19,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [20,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [21,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [22,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [23,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [24,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [25,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [26,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [27,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [28,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [29,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [30,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [31,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "Traceback (most recent call last):\n",
            "  File \"/content/incorrect_index.py\", line 11, in <module>\n",
            "    result = torch.sigmoid(embedding_for_index)\n",
            "             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
            "RuntimeError: CUDA error: device-side assert triggered\n",
            "CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\n",
            "For debugging consider passing CUDA_LAUNCH_BLOCKING=1\n",
            "Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "(though an experienced person might get a hint from the C++ failed assertions printed before the error)\n",
        "\n",
        "There are two ways to find the actual line that caused an exception:\n",
        "\n",
        "* First, you can just move everything to CPU. This is a valid approach, but it involves making changes to your code or input arguments and sometimes can be difficult (for example, if the error occurs after a long chain of operations).\n",
        "\n",
        "* Second, you may use the `CUDA_LAUNCH_BLOCKING` environment variable when starting your code. This will force synchronization for all GPU-related operations, making your code slower but allowing to see the exact source of the error."
      ],
      "metadata": {
        "id": "SK-kBXycW0c-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!CUDA_LAUNCH_BLOCKING=1 python incorrect_index.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HF8lRP3_0XY0",
        "outputId": "9306e04f-730d-4c4e-f763-e3b37e3238e4"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [0,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [1,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [2,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [3,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [4,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [5,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [6,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [7,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [8,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [9,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [10,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [11,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [12,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [13,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [14,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [15,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [16,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [17,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [18,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [19,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [20,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [21,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [22,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [23,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [24,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [25,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [26,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [27,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [28,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [29,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [30,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "../aten/src/ATen/native/cuda/Indexing.cu:1255: indexSelectSmallIndex: block: [0,0,0], thread: [31,0,0] Assertion `srcIndex < srcSelectDimSize` failed.\n",
            "Traceback (most recent call last):\n",
            "  File \"/content/incorrect_index.py\", line 9, in <module>\n",
            "    embedding_for_index = embedding(input)\n",
            "                          ^^^^^^^^^^^^^^^^\n",
            "  File \"/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py\", line 1736, in _wrapped_call_impl\n",
            "    return self._call_impl(*args, **kwargs)\n",
            "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
            "  File \"/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py\", line 1747, in _call_impl\n",
            "    return forward_call(*args, **kwargs)\n",
            "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
            "  File \"/usr/local/lib/python3.11/dist-packages/torch/nn/modules/sparse.py\", line 190, in forward\n",
            "    return F.embedding(\n",
            "           ^^^^^^^^^^^^\n",
            "  File \"/usr/local/lib/python3.11/dist-packages/torch/nn/functional.py\", line 2551, in embedding\n",
            "    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)\n",
            "           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
            "RuntimeError: CUDA error: device-side assert triggered\n",
            "Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Precision of floating point operations\n",
        "Let's have a look at this code snippet and its results. In essence, it computes the sixth power of a random matrix (with a fixed seed) on two different devices."
      ],
      "metadata": {
        "id": "zFDw7mEe0xe1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(1337)\n",
        "x = torch.randn(5000, 5000)\n",
        "torch.use_deterministic_algorithms(False)\n",
        "\n",
        "\n",
        "def matrix_power(x):\n",
        "    y = x @ x @ x @ x @ x @ x\n",
        "    return y.sum().item()\n",
        "\n",
        "\n",
        "print(matrix_power(x))\n",
        "print(matrix_power(x.cuda()))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rNzil3C50_yX",
        "outputId": "d808bd72-7930-42a7-f882-c333d649c5d9"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "27654770130944.0\n",
            "27654807879680.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "(in case you are wondering, `torch.use_deterministic_algorithms(True)` won't help)\n",
        "\n",
        "If we do the same using numpy (in two ways), we also get a different result:"
      ],
      "metadata": {
        "id": "sCl7Jn5Q6gm_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(matrix_power(x.numpy()))\n",
        "np.linalg.matrix_power(x.numpy(), 6).sum()"
      ],
      "metadata": {
        "id": "DDNNwEX72_Bp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a1d76a84-1491-47d4-e0ba-3217ee430294"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "27654749159424.0\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "27654760000000.0"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Takeaway: numerical precision of floating point computations can vary between libraries, environments and devices, and from the user side, it is often hard to resolve this issue altogether. Usually, this happens due to a different summation order in code or due to inherent nondeterminism of hardware.\n",
        "\n",
        "However, note that the relative error is small enough, which makes such blatant discrepancies less of a problem in regular deep learning code."
      ],
      "metadata": {
        "id": "mgGAaBl16uka"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CUDA Graphs\n",
        "\n",
        "In some situations, the bottleneck of your code might be not the GPU compute performance and not even the memory bandwidth, but the latency of kernel execution. Launching each kernel takes CPU time, and if the kernels themselves run very fast, these launches can become a problem.\n",
        "\n",
        "Fortunately, [CUDA graphs](https://github.com/pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/) can address this problem. CUDA graphs capture a sequence of specific operations (i.e., kernel launches) that can later be replayed on new inputs with essentially a single kernel launch. Below, you can see an example of using graphs with a function that consists of many fast-running operations:"
      ],
      "metadata": {
        "id": "KE29v0NDvTup"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def slow_function(x):\n",
        "  for _ in range(500):\n",
        "    y = x*2\n",
        "    z = torch.sigmoid(y)\n",
        "    a = z + 5.0\n",
        "    b = torch.nn.functional.relu(a)\n",
        "    x = b\n",
        "  return b\n",
        "\n",
        "# the code below is a modified version of https://pytorch.org/docs/master/notes/cuda.html#cuda-graph-semantics\n",
        "\n",
        "# Placeholder input used for capture\n",
        "static_input = torch.ones((5,), device=\"cuda\")\n",
        "print(slow_function(static_input))"
      ],
      "metadata": {
        "id": "LEq0c-w0n4Ze",
        "outputId": "7d6c9e48-2677-48a2-8de0-be163e8cd456",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([6.0000, 6.0000, 6.0000, 6.0000, 6.0000], device='cuda:0')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Below, you can see two examples of running complex_function: directly (the standard way) or by capturing `func_graph`, copying new inputs into `static_input` and `replay()`ing it. As you can see, the second version can be alsmost 10x faster compared to default eager execution."
      ],
      "metadata": {
        "id": "1dDnX3EKPiGb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "slow_function(static_input)\n",
        "torch.cuda.synchronize()"
      ],
      "metadata": {
        "id": "k1HOYoQQoULd",
        "outputId": "113b33b6-ebca-4698-9042-72d380a9d386",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 45,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "20.1 ms ± 2.64 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "func_graph = torch.cuda.CUDAGraph()\n",
        "with torch.cuda.graph(func_graph):\n",
        "  static_output = slow_function(static_input)\n",
        "static_output.zero_()\n",
        "\n",
        "# Fills the graph's input memory with new data to compute on\n",
        "static_input.copy_(torch.full((5,), 3, device=\"cuda\"))\n",
        "func_graph.replay()\n",
        "print(static_output)"
      ],
      "metadata": {
        "id": "GGIpgzjgoma-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "be3640c4-3dae-4013-bab0-c3d890b7fde4"
      },
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([6.0000, 6.0000, 6.0000, 6.0000, 6.0000], device='cuda:0')\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "func_graph.replay()\n",
        "torch.cuda.synchronize()"
      ],
      "metadata": {
        "id": "IBD4zgNRo1Qg",
        "outputId": "ff8c4721-fa34-4835-f6ff-258a02aea7d3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.7 ms ± 6.05 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "9UtSWLdqVeoM"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}